# 该文件并不在本机上运行，而是在有GPU的服务器上运行
import os
import random
import re
# from nlpcda.tools.Simbert import Simbert

'''
- 交集复杂问句：询问两个实体之间的相同点（实体数≥2）

| 名称         | 举例                                             |
| ------------ | ------------------------------------------------ |
| 疾病相同症状 | “毛囊炎和银屑病有什么相同的症状吗？”             |
| 症状相同疾病 | “同时具备脓疱和出血性丘疹这些症状的疾病是什么？” |

- 比较句（实体数≥2）
| 名称             | 举例                                                         |
| ---------------- | ------------------------------------------------------------ |
| 疾病治愈率比较   | “乳腺癌和心脏病哪个更容易治疗？”/“乳腺癌和心脏病哪个更难治疗？”/“乳腺癌和心脏病治愈率更大/小？” |
| 疾病治疗周期比较 | "乳腺癌和心脏病哪个治疗周期更长"/“乳腺癌和心脏病哪个要治得更久？” |
'''

# 单位样本扩充数量（n表示一个样本扩充n个相似样本）
num_of_expand = 1


class CQ_Generator:
    def __init__(self):
        cur_dir = '/'.join(os.path.abspath(__file__).split('/')[:-1])
        # 特征词路径
        disease_path = os.path.join(cur_dir, '../../../dict/disease.txt')
        symptom_path = os.path.join(cur_dir, '../../../dict/symptom.txt')
        food_path = os.path.join(cur_dir, '../../../dict/food.txt')
        department_path = os.path.join(cur_dir, '../../../dict/department.txt')
        check_path = os.path.join(cur_dir, '../../../dict/check.txt')
        drug_path = os.path.join(cur_dir, '../../../dict/drug.txt')
        # 加载特征词
        self.disease_wds = [i.strip() for i in open(disease_path, encoding='UTF-8') if i.strip()]
        self.symptom_wds = [i.strip() for i in open(symptom_path, encoding='UTF-8') if i.strip()]
        self.food_wds = [i.strip() for i in open(food_path, encoding='UTF-8') if i.strip()]
        self.department_wds = [i.strip() for i in open(department_path, encoding='UTF-8') if i.strip()]
        self.check_wds = [i.strip() for i in open(check_path, encoding='UTF-8') if i.strip()]
        self.drug_wds = [i.strip() for i in open(drug_path, encoding='UTF-8') if i.strip()]

    # 原始数据生成
    def primary_data_gen(self, label, entities, templates, max_size=100):
        '''
        @param label: 样本标签
        @param entities: 实体类列表
        @param templates: 问句模板
        @param max_size: 数据容量上限
        @return: 生成的自然语言问句
        '''
        res = []
        pattern = r'({.*?})'
        counter = 0  # 生成样本计数器
        while True:
            for temp in templates:
                # 随机生成指定类型的实体
                slots = re.findall(pattern, temp)
                candidates = {}
                for i in range(len(slots)):
                    entity = ''
                    while True:
                        entity = random.choice(entities)
                        if entity not in candidates.values():
                            break
                    candidates[i] = entity
                text = temp.format(candidates)
                res.append({'label': label, 'text': text})
                counter += 1
                if counter >= max_size:
                    break
            print(counter)
            if counter >= max_size:
                break


        return res[:max_size]

    # 数据增强函数
    def data_enhancement(self, primary_data, simbert, enhance_num=1):
        '''
        @param simbert: 预训练数据增强模型
        @param enhance_num: 扩充数目（1代表每个原始样本扩充一条数据）
        @return: 增强后的新数据
        '''
        res = []
        # label = primary_data[0]['label']
        for item in primary_data:
            text = item['text']
            label = item['label']
            synonyms = simbert.replace(sent=text, create_num=enhance_num)
            for item in synonyms:
                res.append({'label': label, 'text': item[0]})

        return res

    # 数据格式化保存（包含数据打乱）
    def data_save(self, dicts, savepath):
        # 打乱数据
        random.shuffle(dicts)
        print("data length: {0}".format(len(dicts)))
        # 保存数据
        fileObject = open(savepath, 'w')
        fileObject.write('label\ttext\n')
        for dictionary in dicts:
            fileObject.write(dictionary['label'] + '\t' + dictionary['text'] + '\n')

        print("---------- Object saved to '{}'".format(savepath))
        fileObject.close()


# 句子生成
def gen_primary_data(gen):
    each_negative_num = 10
    each_positive_num = 200
    positive_samples = []
    negative_samples = []

    '''
            句子模板
        '''
    # 正例部分

    # 疾病相同症状
    intersection_symptom_templates = ["{0[0]}和{0[1]}有什么相同的症状吗？", "{0[0]}与{0[1]}之间有什么一样的症状", "{0[0]}、{0[1]}有什么相同的表现吗",
                                      "{0[0]}、{0[1]}在症状上有什么相同的地方吗？", "{0[0]}、{0[1]}以及{0[2]}在症状上有什么相同点吗？",
                                      "{0[0]}与{0[1]}在症状上有啥交集咩？", "{0[0]}和{0[1]}有什么类似的临床表现？",
                                      "{0[0]},{0[1]}和{0[2]}的症状有啥共通点吗？", "{0[0]}和{0[1]}有什么相似的症状吗？",
                                      "{0[0]}和{0[1]}共同的症状是啥？"]

    # 症状相同疾病
    intersection_disease_templates = ["同时具备{0[0]}和{0[1]}这些症状的疾病是什么？", "什么疾病既包含{0[0]}又有{0[1]}的表现",
                                      "同时有{0[0]}、{0[1]}和{0[2]}等临床表现的疾病有什么？", "在临床上既有{0[0]}同时又有{0[1]}表现的疾病是啥？",
                                      "啥疾病同时有{0[0]}、{0[1]}的症状", "哪些具有{0[0]}表现的疾病同时又有{0[1]}的症状"]

    # 疾病治愈率比较
    compare_disease_cured_prob_templates = ["{0[0]}和{0[1]}哪个更容易治好？", "{0[0]}和{0[1]}哪个更不容易治疗？", "{0[0]}和{0[1]}之间哪一个更难治？",
                                            "{0[0]}和{0[1]}相比哪个治愈率更高？", "{0[0]}和{0[1]}相比哪个治愈率更低？",
                                            "{0[0]}和{0[1]}相比哪个死亡率更高？",
                                            "{0[0]}相比{0[1]}哪个更难痊愈？"]

    # 疾病治疗周期比较
    compare_disease_lasttime_templates = ["{0[0]}与{0[1]}相比哪个治疗周期更长?", "{0[0]}和{0[1]}哪个要治得更久？",
                                          "{0[0]}病人和{0[1]}病人，哪个能更快痊愈？",
                                          "{0[0]}和{0[1]}哪个好的快些呢？", "要花更长时间治疗的疾病是{0[0]}还是{0[1]}？",
                                          "{0[0]}与{0[1]}哪个可以尽快康复？",
                                          "{0[0]}是否比{0[1]}要耗费更长的时间康复？", "{0[0]}是不是比{0[1]}好得慢？"]

    # 正例生成
    positive_samples.extend(
        gen.primary_data_gen("intersection_symptom", gen.disease_wds, intersection_symptom_templates,
                             each_positive_num))
    print('intersection_symptom生成完毕')
    positive_samples.extend(
        gen.primary_data_gen("intersection_disease", gen.symptom_wds, intersection_disease_templates,
                             each_positive_num))
    print('intersection_disease生成完毕')
    positive_samples.extend(
        gen.primary_data_gen("compare_disease_cured_prob", gen.disease_wds, compare_disease_cured_prob_templates,
                             each_positive_num))
    print('compare_disease_cured_prob生成完毕')
    positive_samples.extend(
        gen.primary_data_gen("compare_disease_lasttime", gen.disease_wds, compare_disease_lasttime_templates,
                             each_positive_num))
    print('compare_disease_lasttime生成完毕')
    '''
        反例（简单并列问句）
    '''
    label = "other"

    # 多疾病实体问症状
    negative_temp_1 = ["{0[0]}、{0[1]}和{0[2]}的症状有啥？", "{0[0]}和{0[1]}有啥症状？", "{0[0]}和{0[1]}的症状是？"]
    negative_samples.extend(gen.primary_data_gen(label, gen.disease_wds, negative_temp_1, each_negative_num))

    # 多症状实体问疾病
    negative_temp_2 = ["什么病会有{0[0]}或{0[1]}的症状？", "具有{0[0]}、{0[1]}症状的病是啥？", "啥病有{0[0]}或{0[1]}的临床表现"]
    negative_samples.extend(gen.primary_data_gen(label, gen.symptom_wds, negative_temp_2, each_negative_num))

    # 多疾病实体问成因
    negative_temp_3 = ["{0[0]}和{0[1]}的成因", "为啥有人会得{0[0]}或{0[1]}", "造成{0[0]}或{0[1]}的原因是啥？", "{0[0]}、{0[1]}的原因是什么？"]
    negative_samples.extend(gen.primary_data_gen(label, gen.disease_wds, negative_temp_3, each_negative_num))

    # 多疾病实体问预防措施
    negative_temp_4 = ["{0[0]}和{0[1]}的预防措施", "如何预防{0[0]}或{0[1]}", "怎样才能不得{0[0]}和{0[1]}", "如何才可以不感染{0[0]}以及{0[1]}？"]
    negative_samples.extend(gen.primary_data_gen(label, gen.disease_wds, negative_temp_4, each_negative_num))

    # 多疾病实体问持续时间
    negative_temp_5 = ["{0[0]}和{0[1]}会持续多久", "{0[0]}和{0[1]}要多久才能治好", "{0[0]}、{0[1]}以及{0[2]}的痊愈时间是？",
                       "{0[0]}以及{0[1]}的康复需要多长时间？"]
    negative_samples.extend(gen.primary_data_gen(label, gen.disease_wds, negative_temp_5, each_negative_num))

    # 多疾病实体问治疗方案
    negative_temp_6 = ["{0[0]}和{0[1]}的治疗方案？", "如何治疗{0[0]}和{0[1]}", "{0[0]}和{0[1]}的诊治手段？",
                       "治疗{0[0]}、{0[1]}以及{0[2]}的方法有啥？"]
    negative_samples.extend(gen.primary_data_gen(label, gen.disease_wds, negative_temp_6, each_negative_num))

    # 多疾病实体问治愈概率
    negative_temp_7 = ["{0[0]}和{0[1]}有多大几率能治好？", "{0[0]}和{0[1]}的治愈概率", "{0[0]},{0[1]}和{0[2]}有多大把握康复？"]
    negative_samples.extend(gen.primary_data_gen(label, gen.disease_wds, negative_temp_7, each_negative_num))

    # 多疾病实体问易感人群
    negative_temp_8 = ["{0[0]}和{0[1]}的易感人群是？", "什么人容易感染{0[0]}和{0[1]}？"]
    negative_samples.extend(gen.primary_data_gen(label, gen.disease_wds, negative_temp_8, each_negative_num))

    # 多疾病实体问描述
    negative_temp_9 = ["描述下{0[0]}和{0[1]}？", "什么是{0[0]}和{0[1]}？", "{0[0]}和{0[1]}？"]
    negative_samples.extend(gen.primary_data_gen(label, gen.disease_wds, negative_temp_9, each_negative_num))

    # 多疾病实体问并发症
    negative_temp_10 = ["{0[0]}和{0[1]}的并发症是啥？", "{0[0]}和{0[1]}有啥并发症吗？", "和{0[0]}或{0[1]}并发的疾病是啥？"]
    negative_samples.extend(gen.primary_data_gen(label, gen.disease_wds, negative_temp_10, each_negative_num))

    # 多疾病实体问忌食食物
    negative_temp_11 = ["{0[0]}和{0[1]}的忌食食物有啥？", "{0[0]}和{0[1]}不能吃啥？", "{0[0]}病人和{0[1]}病人需要忌口的食品有什么？"]
    negative_samples.extend(gen.primary_data_gen(label, gen.disease_wds, negative_temp_11, each_negative_num))

    # 多疾病实体问宜食食物
    negative_temp_12 = ["{0[0]}和{0[1]}病人吃什么有好处？", "{0[0]}和{0[1]}病人的宜食食物是啥", "{0[0]}和{0[1]}的推荐食物是什么？",
                        "{0[0]}和{0[1]}患者推荐吃啥？"]
    negative_samples.extend(gen.primary_data_gen(label, gen.disease_wds, negative_temp_12, each_negative_num))

    # 多宜食食物实体问疾病
    negative_temp_13 = ["{0[0]}和{0[1]}对什么病人有好处？", "{0[0]}和{0[1]}适合什么病人食用？", "什么病人推荐吃{0[0]}、{0[1]}和{0[2]}"]
    negative_samples.extend(gen.primary_data_gen(label, gen.food_wds, negative_temp_13, each_negative_num))

    # 多忌食食物实体问疾病
    negative_temp_14 = ["{0[0]}和{0[1]}对什么病人有坏处？", "{0[0]}和{0[1]}不适合什么病人吃？", "什么病食用{0[0]}和{0[1]}不利于康复"]
    negative_samples.extend(gen.primary_data_gen(label, gen.food_wds, negative_temp_14, each_negative_num))

    # 多疾病实体问药物
    negative_temp_15 = ["{0[0]}和{0[1]}的治疗药物是啥？", "{0[0]}和{0[1]}的推荐药品", "什么药品能够治疗{0[0]}、{0[1]}？"]
    negative_samples.extend(gen.primary_data_gen(label, gen.disease_wds, negative_temp_15, each_negative_num))

    # 多药品实体问疾病
    negative_temp_16 = ["{0[0]}和{0[1]}主治什么病？", "{0[0]}以及{0[1]}治疗的病症是啥", "{0[0]}以及{0[1]}能治啥病？"]
    negative_samples.extend(gen.primary_data_gen(label, gen.drug_wds, negative_temp_16, each_negative_num))

    # 多药品实体问生产厂家
    negative_temp_17 = ["{0[0]}和{0[1]}的生产厂家是啥？", "{0[0]}以及{0[1]}在市面上的在售药物是？", "{0[0]}以及{0[1]}有什么商品？"]
    negative_samples.extend(gen.primary_data_gen(label, gen.drug_wds, negative_temp_17, each_negative_num))

    # 多疾病实体问检查
    negative_temp_18 = ["{0[0]}和{0[1]}要做什么检查？", "做什么项目能检查出{0[0]}和{0[1]}？"]
    negative_samples.extend(gen.primary_data_gen(label, gen.disease_wds, negative_temp_18, each_negative_num))

    # 多疾病实体问科室
    negative_temp_19 = ["{0[0]}和{0[1]}的主治科室是啥？", "啥科能诊治{0[0]}和{0[1]}？"]
    negative_samples.extend(gen.primary_data_gen(label, gen.disease_wds, negative_temp_19, each_negative_num))

    # 多检查实体问疾病
    negative_temp_20 = ["{0[0]}和{0[1]}能检查出啥病？", "什么病能通过{0[0]}、{0[1]}查出来？"]
    negative_samples.extend(gen.primary_data_gen(label, gen.check_wds, negative_temp_20, each_negative_num))

    # 多科室实体问疾病
    negative_temp_21 = ["{0[0]}和{0[1]}主治什么疾病", "什么病是去{0[0]}、{0[1]}检查？", "{0[0]}、{0[1]}治疗的病有啥？"]
    negative_samples.extend(gen.primary_data_gen(label, gen.department_wds, negative_temp_21, each_negative_num))
    print("反例生成完毕")
    print("positive sample size: ", len(positive_samples))
    print("negative sample size: ", len(negative_samples))

    return positive_samples, negative_samples


if __name__ == '__main__':
    # 初始化生成器
    gen = CQ_Generator()
    # Simbert模型
    # config = {
    #     'model_path': '/home/cike/diploma project/code_reference/PaddleHub/demo/my_relation_classification/utils/nlpcda/tools/models/chinese_simbert_L-12_H-768_A-12',
    #     'CUDA_VISIBLE_DEVICES': '6',
    #     'max_len': 64,
    #     'seed': 1
    # }
    # simbert = Simbert(config=config)
    # 生成原始数据
    positive_samples, negative_samples = gen_primary_data(gen)

    # positive_enhance = gen.data_enhancement(positive_samples,simbert)

    data = []
    data.extend(positive_samples)
    data.extend(negative_samples)
    # data.extend(positive_enhance)
    print("----- length of generated data: {0} -----".format(len(data)))
    random.shuffle(data)
    train_length = int(len(data)*0.6)
    dev_length = int(len(data)*0.2)
    train = data[:train_length]
    dev = data[train_length:dev_length+train_length]
    test = data[dev_length+train_length:]
    # 数据保存
    # gen.data_save(data, "./ComplexQ_data/primary_data.txt")
    gen.data_save(train, "./ComplexQ_data/train.txt")
    gen.data_save(dev, "./ComplexQ_data/dev.txt")
    gen.data_save(test, "./ComplexQ_data/test.txt")


